import json
import numpy as np
import os
import matplotlib.pyplot as plt
from collections import Counter

#初始化一个保存类别信息的列表
key_class_list = []
key_class_truncated_list = []
key_class_sheltered_list = []

def objetct_class(files):
    #读取标注信息并写入 xml
    for json_file_ in files:
        json_filename = os.path.join(labelme_path, json_file_) + ".json"
        json_file = json.load(open(json_filename, "r", encoding="gb2312"))
        
        for classes in json_file["shapes"]:
            label = classes["label"]
            key_class_list.append((label))
            count = Counter(key_class_list)
            classlabel = list(count)
            classlabel=np.array(classlabel)
            num = count.values()
    mean = sum(num) / len(files)
    # print("mean bbox of dataset is %s"% mean)

    #对类别信息与类别数目进行可视化

    #刻度距离坐标轴的距离调整
    plt.tick_params(pad = 0.03)  #通过pad参数调整距离

    plt.barh(classlabel, num, color = 'g', align = 'center') 
    plt.text(max(num)/2, len(num)*1.1, list(num), ha='center', fontsize=12)
    plt.text(max(num)/2, len(num)*1.2, "imgs of dataset is %s"% len(files), ha='center', fontsize=12)
    plt.text(max(num)/2, len(num)*1.3, "mean bbox of dataset is %s"% mean, ha='center', fontsize=12)
    plt.title('Bosch 2D Detection') 
    plt.ylabel('Class axis',fontsize=11) 
    plt.xlabel('Num axis',fontsize=11,linespacing = 0.3, labelpad = 0.3) 
    plt.savefig('class_num.png', bbox_inches='tight')
    plt.close()

def objetct_truncated(files):
    #读取标注信息并写入 xml
    for json_file_ in files:
        json_filename = os.path.join(labelme_path, json_file_) + ".json"
        json_file = json.load(open(json_filename, "r", encoding="gb2312"))
        
        for classes in json_file["shapes"]:
            truncated = classes["truncated"]
            key_class_truncated_list.append((truncated))
            count_truncated = Counter(key_class_truncated_list)
            class_truncated = list(count_truncated)
            class_truncated=np.array(class_truncated)
            num_class_truncated = count_truncated.values()

    #对类别信息与类别数目进行可视化

    #刻度距离坐标轴的距离调整
    plt.tick_params(pad = 0.03)  #通过pad参数调整距离

    plt.barh(class_truncated, num_class_truncated, color = 'g', align = 'center') 
    plt.text(max(num_class_truncated)/2, len(num_class_truncated)*1.05, "imgs of dataset is %s"% len(files), ha='center', fontsize=12)
    plt.text(max(num_class_truncated)/2, len(num_class_truncated)*1.1, list(num_class_truncated), ha='center', fontsize=12)
    plt.text(max(num_class_truncated)/2, len(num_class_truncated)*1.25, '0: not truncated', ha='center', fontsize=12)
    plt.text(max(num_class_truncated)/2, len(num_class_truncated)*1.3, '1: horizon truncated', ha='center', fontsize=12)
    plt.text(max(num_class_truncated)/2, len(num_class_truncated)*1.35, '2: vertical truncated', ha='center', fontsize=12)
    plt.text(max(num_class_truncated)/2, len(num_class_truncated)*1.4, '3: horizon + vertical truncated', ha='center', fontsize=12)
    plt.title('Bosch 2D Detection') 
    plt.ylabel('Truncated axis',fontsize=11) 
    plt.xlabel('Num axis',fontsize=11,linespacing = 0.3, labelpad = 0.3) 
    plt.savefig('class_truncated.png', bbox_inches='tight')
    plt.close()

def objetct_sheltered(files):
    #读取标注信息并写入 xml
    for json_file_ in files:
        json_filename = os.path.join(labelme_path, json_file_) + ".json"
        json_file = json.load(open(json_filename, "r", encoding="gb2312"))
        
        for classes in json_file["shapes"]:
            sheltered = classes["sheltered"]
            key_class_sheltered_list.append((sheltered))
            count_sheltered = Counter(key_class_sheltered_list)
            class_sheltered = list(count_sheltered)
            class_sheltered=np.array(class_sheltered)
            num_class_sheltered = count_sheltered.values()

    #对类别信息与类别数目进行可视化

    #刻度距离坐标轴的距离调整
    plt.tick_params(pad = 0.03)  #通过pad参数调整距离

    plt.barh(class_sheltered, num_class_sheltered, color = 'g', align = 'center') 
    plt.text(max(num_class_sheltered)/2, len(num_class_sheltered)*1.05, "imgs of dataset is %s"% len(files), ha='center', fontsize=12)
    plt.text(max(num_class_sheltered)/2, len(num_class_sheltered)*1.1, list(num_class_sheltered), ha='center', fontsize=12)
    plt.text(max(num_class_sheltered)/2, len(num_class_sheltered)*1.25, '0: not sheltered', ha='center', fontsize=12)
    plt.text(max(num_class_sheltered)/2, len(num_class_sheltered)*1.3, '1: 0-50 sheltered', ha='center', fontsize=12)
    plt.text(max(num_class_sheltered)/2, len(num_class_sheltered)*1.35, '2: 50-80 sheltered', ha='center', fontsize=12)
    plt.text(max(num_class_sheltered)/2, len(num_class_sheltered)*1.4, '3: 80-100 sheltered', ha='center', fontsize=12)
    plt.title('Bosch 2D Detection') 
    plt.ylabel('Sheltered axis',fontsize=11) 
    plt.xlabel('Num axis',fontsize=11,linespacing = 0.3, labelpad = 0.3) 
    plt.savefig('class_sheltered.png', bbox_inches='tight')
    plt.close()


if __name__ == "__main__":
    #标签路径
    labelme_path = "all_json_104806_20220729"   #原始labelme标注数据路径
    #获取待处理文件
    files = [p.split('.json')[0] for p in os.listdir(labelme_path) if '.json' in p]
    print("there are %s files"%len(files))
    objetct_class(files)
    objetct_truncated(files)
    objetct_sheltered(files)